from gear import HiddenPrints, PrintToTxt, PlanetaryGear, PlanetInit, Gear, GearInit

from math import sin, cos, tan, radians, degrees, floor, sqrt, acos, atan, pi, log
from invfun import inv, inv_1
import numpy as np
from matplotlib import pyplot as plt
import os


def deltafun_sc(gear: Gear, select: int, alpha_x, B: float, F=None):
    alpha = gear.alpha
    if gear.InitStruct.types == 'wai':
        typeflag = 1
    else:  # 内啮合副，要判断哪个是内齿轮（一般下标2是内齿轮，但还是判断一下）
        zflag = int(gear.z2 > gear.z1)  # 2是内齿轮,为1；反之为-1
        if select == 1:
            typeflag = 1*zflag
        elif select == 2:
            typeflag = -1*zflag

    if select == 1:
        #material = gear.InitStruct.material1
        E = gear.InitStruct.E1
        v = gear.InitStruct.v1
        db = gear.db1
        s = gear.s1
        d = gear.d1
        dNf = gear.dNf1
        #z = gear.z1
        df = gear.df1
        da = gear.da1
        alpha_a = gear.alpha_a1
        #x = gear.x1

    elif select == 2:
        #material = gear.InitStruct.material2
        E = gear.InitStruct.E2
        v = gear.InitStruct.v2
        db = gear.db2
        s = gear.s2
        d = gear.d2
        dNf = gear.dNf2
        #z = gear.z2
        df = gear.df2
        da = gear.da2
        alpha_a = gear.alpha_a2
        #x = gear.x2
    else:
        return

    r_x = (db/2)/cos(alpha_x)  # 啮合点所在从动轮圆半径

    # s_x = s*r_x/(d/2)+2*r_x * (inv(alpha)-inv(alpha_x))  # 啮合点齿厚
    # gama_x = s_x/(2*r_x)  # 啮合点半齿厚对应的圆心角
    gama_x = s/d + (inv(alpha) - inv(alpha_x))*typeflag
    beta_x = alpha_x - gama_x  # 载荷作用角

    # x_j = r_x*cos(gama_x)  # 载荷作用点坐标
    #y_j = r_x*sin(gama_x)
    # 设基点在齿根圆
    if dNf >= db:  # 渐开线起始圆直径大于基圆直径
        r_lj = dNf/2  # 过渡圆角临界半径
    else:
        r_lj = db/2

    alpha_lj = acos(db/2/r_lj)  # 临界半径的压力角
    s_lj = 2*r_lj*sin(
        s/d +
        (inv(alpha)-inv(alpha_lj))*typeflag
    )  # 临界半径的齿厚
    # 以下是石川公式的下标，含义不再做说明
    #s_f = df*sin(s/d+inv(alpha))
    tmp = sqrt((df/2)**2-(s_lj/2)**2)
    h_r = (sqrt(r_lj**2-(s_lj/2)**2) - tmp)*typeflag
    h_x = (r_x*cos(gama_x) - tmp)*typeflag
    s_a = da*sin(s/d +
                 (inv(alpha)-inv(alpha_a))*typeflag
                 )
    h = (sqrt((da/2)**2-(s_a/2)**2) - tmp)*typeflag
    h_i = (h*s_lj-h_r*s_a)/(s_lj-s_a)
    # 变形量用delta表示，不算F
    delta_br = 12*(cos(beta_x)**2)/(E*B*(s_lj**3)) * (
        h_x*h_r*(h_x-h_r)+(h_r**3)/3
    )
    tmp2 = (h_i-h_x)/(h_i-h_r)

    delta_bt = 6*(cos(beta_x)**2)/(E*B*(s_lj**3)) * (
        tmp2*(4 - tmp2) - 2*log(tmp2) - 3
    )*((h_i-h_r)**3)

    delta_s = 2*(1+v)*(cos(beta_x)**2)/(E*B*s_lj)*(
        h_r - (h_i-h_r)*log(tmp2)
    )  # 稍微修改了公式的写法，+变-，log里面取倒数

    delta_g = 24*(h_x**2)*(cos(beta_x)**2)/(pi*E*s_lj**2*B)
    deltaall = delta_br+delta_bt+delta_g+delta_s
    return deltaall


def shichuan(gear: Gear, x, B, F=None):  # B 齿宽

    typeflag = 1 if gear.InitStruct.types == 'wai' else -1
    '''
    alpha = gear.alpha  # 压力角
    # 材料参数
    material2 = gear.InitStruct.material2
    E2 = gear.InitStruct.E2
    v2 = gear.InitStruct.v2

    material1 = gear.InitStruct.material1
    E1 = gear.InitStruct.E1
    v1 = gear.InitStruct.v1
    '''
    # x：以从动轮齿顶圆开始接触算起的转角，设从动轮是齿轮2

    # 计算从动轮2的一对齿啮合的x最大值
    n1n2 = gear.a_work*sin(gear.alpha_work)
    x_max = (tan(gear.alpha_a2) -
             (n1n2-typeflag*sqrt((gear.da1/2)**2-(gear.db1/2)**2))/(gear.db2/2))*typeflag
    # print(x_max)
    x = x % x_max
    alpha_x2 = atan(
        tan(gear.alpha_a2) - typeflag * x
    )  # 转动x度时啮合点在从动轮齿上的压力角

    delta_2 = deltafun_sc(gear, 2, alpha_x2, B, F)  # 齿轮2的变形

    alpha_x1 = atan(
        (n1n2 - (gear.db2/2)*tan(alpha_x2))*typeflag / (gear.db1/2)
    )  # 转动x度时啮合点在主动轮齿上的压力角
    delta_1 = deltafun_sc(gear, 1, alpha_x1, B, F)  # 齿轮1的变形
    delta_pv = 4*(1+pi/4)/pi/B*(
        (1-gear.InitStruct.v1**2)/gear.InitStruct.E1 +
        (1-gear.InitStruct.v2**2)/gear.InitStruct.E2
    )

    deltaall = delta_1+delta_2+delta_pv
    if F != None:
        deltaall = deltaall*F
    return deltaall





def stiffness(gg: Gear, xstep=radians(0.05), savefilename=None, show=False):
    typeflag = 1 if gg.InitStruct.types == 'wai' else -1
    n1n2 = gg.a_work*sin(gg.alpha_work)
    c1 = (n1n2-sqrt((gg.da2/2)**2-(gg.db2/2)**2))*typeflag
    p_b = pi*gg.db2/gg.z2  # 基圆齿距
    c4 = c1 + p_b
    c5 = sqrt((gg.da1/2)**2-(gg.db1/2)**2)
    c2 = c5 - p_b
    tan_alpha_x_max = (n1n2-typeflag*c5)/(gg.db2/2)
    x_max = typeflag*(tan(gg.alpha_a2) - tan_alpha_x_max)  # 单齿最大转角
    tan_alpha_x_2to1 = (n1n2-typeflag*c2)/(gg.db2/2)
    x_2to1 = typeflag*(tan(gg.alpha_a2) - tan_alpha_x_2to1)  # 双齿啮合取进单齿啮合区
    tan_alpha_x_1to2 = (n1n2-typeflag*c4)/(gg.db2/2)
    x_1to2 = typeflag*(tan(gg.alpha_a2) - tan_alpha_x_1to2)  # 单齿啮合区进双齿啮合区

    xrange1 = np.arange(0, (x_2to1), xstep)
    xrange2 = np.arange((x_2to1), (x_1to2), xstep)
    xrange3 = np.arange((x_1to2), (x_max), xstep)
    xrange = np.concatenate((xrange1, xrange2, xrange3))
    stiff_1 = []
    stiff_2_1 = []
    stiff_2_3 = []
    B = 40
    #T = 90
    #F = T/(gg.db1/2)
    stiff_all = []
    stiff_max = 0  # 啮合过程中的刚度最值
    stiff_min = 1e5

    for x in xrange1:  # 双齿啮合区1
        tan_alpha_x = tan(gg.alpha_a2) - typeflag * x
        tan_alpha_x_another = (gg.db2/2*tan_alpha_x - typeflag*p_b)/(gg.db2/2)
        x_another = typeflag*(tan(gg.alpha_a2) - tan_alpha_x_another)
        st1 = 1/(shichuan(gg, x, B)*B)
        st2 = 1/(shichuan(gg, x_another, B)*B)
        st_all = st1+st2
        if stiff_max < st_all:
            stiff_max = st_all
        stiff_all.append(st_all)  # 总刚度
        stiff_1.append(st1)  # 单齿的刚度（刚进啮合区）
        stiff_2_1.append(st2)  # 单齿的刚度（马上离开啮合区）
    st_last = st_all  # 保存最后的一个刚度值，用于计算第一次跳变
    tmpflag = True

    for x in xrange2:  # 单齿啮合区
        st1 = 1/(shichuan(gg, x, B)*B)
        if stiff_min > st1:
            stiff_min = st1
        if tmpflag:
            st_jump_1 = st_last - st1  # 计算第一次跳变，仅计算一次
            tmpflag = False
        stiff_all.append(st1)  # 总刚度
        stiff_1.append(st1)  # 单齿的刚度
    st_last = st1  # 保存最后的一个刚度值，用于计算第二次跳变
    tmpflag = True

    for x in xrange3:  # 双齿啮合区
        tan_alpha_x = tan(gg.alpha_a2) - typeflag * x
        tan_alpha_x_another = (gg.db2/2*tan_alpha_x + typeflag*p_b)/(gg.db2/2)
        x_another = typeflag*(tan(gg.alpha_a2) - tan_alpha_x_another)
        st1 = 1/(shichuan(gg, x, B)*B)
        st2 = 1/(shichuan(gg, x_another, B)*B)
        st_all = st1+st2
        if stiff_max < st_all:
            stiff_max = st_all
        if tmpflag:
            st_jump_2 = st_all - st_last  # 计算第一次跳变，仅计算一次
            tmpflag = False
        stiff_all.append(st_all)  # 总刚度
        stiff_1.append(st1)  # 单齿的刚度（马上离开啮合区）（所研究的目标齿轮）
        stiff_2_3.append(st2)  # 单齿的刚度（刚进啮合区）

    if (savefilename != None) or show:  # 如果既不储存也不显示，就不作图，仅仅返数据
        plt.figure(figsize=(5, 5))
        plt.plot(xrange*180/pi, stiff_all, 'r-', label='Double Stiffness')
        plt.plot(xrange*180/pi, stiff_1, 'b--',label='Single(this) Stiffness')
        
        plt.plot(xrange1*180/pi, stiff_2_1, 'g--',label='Single(prev) Stiffness')

        plt.plot(xrange3*180/pi, stiff_2_3, 'g--',label='Single(next) Stiffness')

        plt.xlabel('Angle Rotate/(deg)')
        plt.ylabel('Contact Stiffness/(N/(mm${\\cdot\\mu}$m))')
        plt.title('Stiffness -- Rotate Angle')
        plt.grid()
        plt.legend()
        if savefilename != None:
            if not os.path.exists('./pic'):
                os.mkdir('./pic')
            plt.savefig("./pic/"+savefilename, dpi=300, bbox_inches='tight')
        if show:
            plt.show()

    return ((xrange, stiff_all, stiff_1), (xrange1, stiff_2_1), (xrange3, stiff_2_3)), (stiff_max, stiff_min, max(st_jump_1, st_jump_2))


def pt_stiffness(pt: PlanetaryGear, xstep=radians(0.05), savefilename=None, show=False):
    xs1, maxmin1 = stiffness(gg=pt.gear_ac, xstep=xstep,
                             savefilename=None, show=False)
    # 不在stiffness()中进行绘图，在本函数中获取数据进行一张图的绘制
    xs2, maxmin2 = stiffness(gg=pt.gear_cb, xstep=xstep,
                             savefilename=None, show=False)
    if (savefilename != None) or show:  # 如果既不储存也不显示，就不作图，仅仅返数据
        plt.figure(figsize=(18,9))
        
        plt.subplot(1,2,1)
        plt.plot(xs1[0][0]*180/pi, xs1[0][1], 'r-',
                label='Double Stiffness')
        plt.plot(xs1[0][0]*180/pi, xs1[0][2], 'b--',
                label='Single(this) Stiffness')
        #x_range_1and3 = np.append(xrange1, xrange3)
        plt.plot(xs1[1][0]*180/pi, xs1[1][1], 'g--',
                label='Single(prev) Stiffness')

        plt.plot(xs1[2][0]*180/pi, xs1[2][1], 'g--',
                label='Single(next) Stiffness')
        plt.xlabel('Angle Rotate/(deg)')
        plt.ylabel('Contact Stiffness/(N/(mm${\\cdot\\mu}$m))')
        plt.title('sun&planet Stiffness -- Rotate Angle')
        plt.grid()
        plt.legend()
        
        plt.subplot(1,2,2)
        plt.plot(xs2[0][0]*180/pi, xs2[0][1], 'r-',
                label='Double Stiffness')
        plt.plot(xs2[0][0]*180/pi, xs2[0][2], 'b--',
                label='Single(this) Stiffness')
        #x_range_1and3 = np.append(xrange1, xrange3)
        plt.plot(xs2[1][0]*180/pi, xs2[1][1], 'g--',
                label='Single(prev) Stiffness')

        plt.plot(xs2[2][0]*180/pi, xs2[2][1], 'g--',
                label='Single(next) Stiffness')
        plt.xlabel('Angle Rotate/(deg)')
        plt.ylabel('Contact Stiffness/(N/(mm${\\cdot\\mu}$m))')
        plt.title('internal&planet Stiffness -- Rotate Angle')
        plt.grid()
        plt.legend()
        if savefilename != None:
            if not os.path.exists('./pic'):
                os.mkdir('./pic')
            plt.savefig("./pic/"+savefilename, dpi=400, bbox_inches='tight')
        if show:
            plt.show()
            
        
    
    
    pt.gear_ac.st_jump_max = maxmin1[2]
    pt.gear_cb.st_jump_max = maxmin2[2]
    return (xs1,xs2),(maxmin1,maxmin2)

if __name__ == '__main__':
    
    init1 = PlanetInit()

    #planet1 = PlanetaryGear(1.75, 18, 44, 108, 0.5863, 0.691, 0.830, init1)
    planet1 = PlanetaryGear(1.75, 18, 44, 108, 0.5863, 0.391, 0.830, init1)
    # planet1.PrintItem()

    g13 = planet1.gear_ac
    g32 = planet1.gear_cb
    

    step = radians(0.01)
    
    stiffness(g13,step,show=True)
    '''
    xs,maxmin = pt_stiffness(planet1,xstep=step,savefilename='shichuanstiff_pt2.png',show=True)
    
    print(maxmin[0])
    print(maxmin[1])
    '''
    
    